import dramatiq
from fastapi import Depends
from dramatiq.brokers.redis import RedisBroker
from app.utils.response_utils import build_response
from app.controllers.upload_controller import apply_file_upload_lease, update_ai_upload_file_record
from app.models.new_upload_model import upload_file
from app.controllers.new_upload_controller import add_file
from app.database.db_setup import get_db
from sqlalchemy.orm import Session
import time
import logging
from contextlib import contextmanager

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 配置 Redis 作为消息代理
redis_broker = RedisBroker(host="localhost", port=6379)
dramatiq.set_broker(redis_broker)


# 定义上下文管理器类用于依赖注入
@contextmanager
def get_db_context():
    db = get_db()
    try:
        yield next(db)  # 获取生成器返回的会话对象
    finally:
        try:
            next(db)  # 关闭会话
        except StopIteration:
            pass


@dramatiq.actor
def process_upload_task(md5, cat_id, category_id, file_name, size_in_bytes, absolute_file_path, id):
    """
    模拟处理上传后的任务，你可以在这里添加实际的业务逻辑
    :param md5: 文件的 MD5 值
    :param cat_id: 分类 ID
    :param category_id: 分类 ID
    :param file_name: 名称
    :param size_in_bytes: 文件大小
    :param add_time: 添加时间
    """
    try:
        # 调用上传文件函数
        logger.info(f"开始调用 apply_file_upload_lease 函数，文件名为 {file_name}")
        lease_result = apply_file_upload_lease(category_id, file_name, size_in_bytes, md5)
        if "error" in lease_result:
            logger.error(f"apply_file_upload_lease 失败，错误信息: {lease_result['error']}")
            return build_response(100, '失败', '名字已存在')
        logger.info(f"apply_file_upload_lease 成功，返回结果: {lease_result}")

        # 从返回结果中获取 body 部分
        body = lease_result.get("body")
        if body is None:
            logger.error("apply_file_upload_lease 返回结果中 body 为空")
            return build_response(100, '失败', "返回结果有误")

        # 从 body 中获取 Data 部分
        data = body.get("Data")
        if data is None:
            logger.error("apply_file_upload_lease 返回结果中 Data 为空")
            return build_response(100, '失败', "返回结果有误")

        FileUploadLeaseId = data.get("FileUploadLeaseId")
        logger.info(f"获取到的 FileUploadLeaseId: {FileUploadLeaseId}")

        if not FileUploadLeaseId:
            logger.error("apply_file_upload_lease 返回结果中 FileUploadLeaseId 为空")
            return build_response(100, '失败', "返回结果有误")
        # 从 body 中获取 Param 部分
        param = data.get("Param")
        if param is None:
            logger.error("apply_file_upload_lease 返回结果中 Param 为空")
            return build_response(100, '失败', "返回结果有误")

        headers = param.get("Headers")
        if headers is None:
            logger.error("apply_file_upload_lease 返回结果中 Headers 为空")
            return build_response(100, '失败', "返回结果有误")
        logger.info(f"获取到的 headers: {headers}")
        extra = headers.get("X-bailian-extra")
        logger.info(f"获取到的 extra: {extra}")
        if not extra:
            logger.error("apply_file_upload_lease 返回结果中 X-bailian-extra 为空")
            return build_response(100, '失败', "返回结果有误")
        content_type = headers.get("Content-Type")
        if not content_type:
            logger.error("apply_file_upload_lease 返回结果中 Content-Type 为空")
            return build_response(100, '失败', "返回结果有误")
        pre_signed_url = param.get("Url")
        if not pre_signed_url:
            logger.error("apply_file_upload_lease 返回结果中 Url 为空")
            return build_response(100, '失败', "返回结果有误")
        logger.info(f"获取到的 pre_signed_url: {pre_signed_url}")
        logger.info(f"获取到的 absolute_file_path: {absolute_file_path}")
        logger.info(f"获取到的 extra: {extra}")
        logger.info(f"获取到的 content_type: {content_type}")

        logger.info(f"开始调用 upload_file 函数，文件名为 {file_name}")
        res = upload_file(pre_signed_url, absolute_file_path, extra, content_type)
        if res:
            logger.info(f"upload_file 成功，文件名为 {file_name}")
            logger.info(f"开始调用 add_file 函数，文件名为 {file_name}")
            add_res = add_file(FileUploadLeaseId, 'DASHSCOPE_DOCMIND', category_id, file_name)
            if "error" in add_res:
                logger.error(f"add_file 失败，错误信息: {add_res['error']}")
                return build_response(100, '失败', '名字已存在')
            logger.info(f"add_file 成功，返回结果: {add_res}")

            # 从返回结果中获取 body 部分
            body = add_res.get("body")
            if body is None:
                logger.error("add_file 返回结果中 body 为空")
                return build_response(100, '失败', "返回结果有误")

            # 从 body 中获取 Data 部分
            data = body.get("Data")
            if data is None:
                logger.error("add_file 返回结果中 Data 为空")
                return build_response(100, '失败', "返回结果有误")
            FileId = data.get("FileId")
            logger.info(f"获取到的 FileId: {FileId}")
            if not FileId:
                logger.error("add_file 返回结果中 FileId 为空")
                return build_response(100, '失败', "返回结果有误")
            Parser = data.get("Parser")
            logger.info(f"获取到的 Parser: {Parser}")
            if not Parser:
                logger.error("add_file 返回结果中 Parser 为空")
                return build_response(100, '失败', "返回结果有误")
            set_time = int(time.time())
            # 使用上下文管理器获取 db 并调用 update_ai_upload_file_record
            with get_db_context() as db:
                logger.info(f"开始调用 update_ai_upload_file_record 函数，文件名为 {file_name}")
                res = update_ai_upload_file_record(db, id, FileId, Parser, 1, set_time)
            logger.info(f"update_ai_upload_file_record 返回参数: {res}")
        else:
            logger.error(f"upload_file 失败，文件名为 {file_name}")
    except Exception as e:
        logger.exception(f"处理上传任务时发生未知错误: {e}")
        return build_response(100, '失败', str(e))